package cc.git.liuyan.customeraiagent.core.embeddingmodel.abstrct;

import cc.git.liuyan.customeraiagent.core.embeddingmodel.*;

import java.math.BigDecimal;
import java.util.ArrayList;
import java.util.List;

public abstract class AbstractHybridEmbeddingModel implements HybridEmbeddingModel {
    protected EmbeddingModel denseEmbeddingModel;
    protected EmbeddingModel spaseEmbeddingModel;
    protected EmbeddingModelSplit split;

    public AbstractHybridEmbeddingModel(EmbeddingModel denseEmbeddingModel, EmbeddingModel spaseEmbeddingModel) {
        this.denseEmbeddingModel = denseEmbeddingModel;
        this.spaseEmbeddingModel = spaseEmbeddingModel;
    }

    @Override
    public final EmbeddingModel getDenseEmbeddingModel() {
        return this.denseEmbeddingModel;
    }

    @Override
    public EmbeddingModel getSpaseEmbeddingModel() {
        return this.spaseEmbeddingModel;
    }

    @Override
    public String modelName() {
        return null;
    }

    @Override
    public int dimension() {
        return 0;
    }

    @Override
    public final void setSplit(EmbeddingModelSplit split) {
        this.split = split;
    }

    @Override
    public final EmbeddingModelSplit getSplit() {
        return this.split;
    }

    @Override
    public final List<EmbeddingModelOutputData> embed(EmbeddingModelInputData inputData) {
        if (this.split == null) {
            this.split = new SimpleEmbeddingModelSplit();
        }
        List<EmbeddingModelOutputData> outputDatas = new ArrayList<EmbeddingModelOutputData>();
        //不管输入是什么类型，分片后segment数据都是type=string
        List<EmbeddingModelInputData> inputDataSegmentList = split.split(inputData);
        for (EmbeddingModelInputData embeddingModelInputData : inputDataSegmentList) {
            EmbeddingModelOutputData outputData = new EmbeddingModelOutputData(embeddingModelInputData.getInputData());
            outputData.setDenseVectors(this.denseEmbeddingModel.embed(embeddingModelInputData.getInputData()));
            outputData.setSparseVectors(this.spaseEmbeddingModel.embed(embeddingModelInputData.getInputData()));
            outputDatas.add(outputData);
        }

        return outputDatas;
    }

    @Override
    public final List<BigDecimal> embed(String text) {
        return null;
    }
}
